#include <stdio.h>
#include <string.h>
#include <cv.h>
#include <cvaux.h>
#include <highgui.h>
#include <math.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <getopt.h>

// ścieżka do pliku detektora
#define DETECTOR "./haarcascade_frontalface_default.xml"

// domyślne ustawienia
int _DEVICE = CV_CAP_ANY; // numer urządzenia
int _PICTURES = 5; // liczba zdjęć
int _Y = 30; // zwiększenie wykrytego obszaru
int _SIZE_Y = 100; // rozmiar zdjęcia po skalowaniu

// zmienne pomocnicze
int _break = 0;
int _skip = 0;
char _DIR[55] = { 0 };

// deklaracje funkcji
IplImage* detect(IplImage* pInpImg);
IplImage* snapshot();
void collect(int count, int id, int start);
void show(IplImage* pImage);
void printUsage();

int main(int argc, char* argv[]) {
	int c;
	int classID;

	if(argc == 2 && (strcmp(argv[1], "help") == 0 || strcmp(argv[1], "--help") == 0)) {
		printUsage();
		return 0;
	}

	while((c = getopt(argc, argv, "c:d:e:p:s:")) != -1) {
		switch(c) {
			case 'c':
				classID = atoi(optarg);
				break;
			case 'd':
				_DEVICE = atoi(optarg);
				break;
			case 'e':
				_Y = atoi(optarg);
				break;
			case 'p':
				_PICTURES = atoi(optarg);
				break;
			case 's':
				_SIZE_Y = atoi(optarg);
				break;
			case '?':
				fprintf(stderr, "Invalid command or no argument specified. Type \"collect --help\" to view valid commands.\n", optopt);
				exit(1);
			default:
				abort();
		}
	}

	sprintf(_DIR, "class%d", classID);
						
	mkdir(_DIR, 0755);
		
	collect(_PICTURES, 1, 1);
	return 0;
}

void printUsage() {
	fprintf(stderr, "Usage: collect <commands>\n"
			"Valid commands are\n"
			"\t-c\tPerson's class identifier. E.g. -c 2\n"
			"\t-d\tCamera device number. Value X of /dev/videoX . E.g. -d 0\n"
			"\t-e\tAdditional pixels for region of interest. E.g. -e 30\n"
			"\t-p\tNumber of pictures. E.g. -p 5\n"
			"\t-s\tSize of the image. E.g. -s 100\n"
		);
}

IplImage* detect(IplImage* pInpImg) {
	CvHaarClassifierCascade * pCascade = 0; 
	CvMemStorage* pStorage = 0;
	CvSeq* pFaceRectSeq;
	IplImage* pFaceImg = 0;
	int i;
	
	pStorage = cvCreateMemStorage(0);
	pCascade = (CvHaarClassifierCascade*)cvLoad(DETECTOR, 0, 0, 0);
	
	if(!pInpImg || !pStorage || !pCascade)
		return 0;

	pFaceRectSeq = cvHaarDetectObjects(pInpImg, pCascade, pStorage, 1.1, 2, CV_HAAR_DO_CANNY_PRUNING, cvSize(60, 60));
		
	if(pFaceRectSeq->total > 0) {
		CvRect* pFaceRect = (CvRect*)cvGetSeqElem(pFaceRectSeq, 0);
		
		// Poniższy kod poprawia kadrowanie twarzy wykrytej przez detektor:
			
		// powiększenie obszaru
		const int w = _Y;
		const int h = _Y;
		pFaceRect->x -= w;
		pFaceRect->y -= h;
		pFaceRect->width += 2*w;
		pFaceRect->height += 2*h;			
		// przesunięcie ROI w górę
		pFaceRect->y -= floor(0.09 * pFaceRect->height);			
		// zmniejszenie szerokości
		pFaceRect->width = 0.8*pFaceRect->width;
		pFaceRect->x += pFaceRect->height*0.1;
		const int _SIZE_X = 0.8*_SIZE_Y;
		
		// Skopiowanie, transformacja do skali odcieni szarości i skalowanie fotografii
		IplImage* pPreview = cvCloneImage(pInpImg);
		CvPoint pt1 = { pFaceRect->x, pFaceRect->y };
		CvPoint pt2 = { pFaceRect->x + pFaceRect->width, pFaceRect->y + pFaceRect->height };
		cvRectangle(pPreview, pt1, pt2, CV_RGB(0,255,0), 3, 4, 0);
		// wyświetlenie podglądu
		show(pPreview);
		
		cvSetImageROI(pInpImg, *pFaceRect);
		IplImage* pColorFaceImg = cvCreateImage(cvSize(_SIZE_X, _SIZE_Y), pInpImg->depth, pInpImg->nChannels);
		pFaceImg = cvCreateImage(cvGetSize(pColorFaceImg), IPL_DEPTH_8U, 1);
		
		cvResize(pInpImg, pColorFaceImg, CV_INTER_AREA);
		cvCvtColor(pColorFaceImg, pFaceImg, CV_BGR2GRAY);
		
		cvResetImageROI(pInpImg);
	}
	
	// czyszczenie
	if(pCascade)
		cvReleaseHaarClassifierCascade(&pCascade);
	
	if(pStorage)
		cvReleaseMemStorage(&pStorage);
		
	return pFaceImg;
}

IplImage* snapshot() {
	int i;
	CvCapture* pCapture = 0;
	IplImage* pVideoFrame = 0;
	IplImage* pSnapShot = 0;
	
	pCapture = cvCaptureFromCAM(_DEVICE);
	if(!pCapture) {
		fprintf(stderr, "Failed to initialize video capture\n");
		return 0;
	}
	
	// przed pobraniem zdjęcia należy wykonać kilka zdjęć
	// na podstawie których kamera ustawia ostrość i kontrast
	for(i = 0; i < 25; ++i)
		pVideoFrame = cvQueryFrame(pCapture);

	pVideoFrame = cvQueryFrame(pCapture);
	if(!pVideoFrame) {
		fprintf(stderr, "Failed to get a video frame\n");
		return 0;
	}

	pSnapShot = cvCloneImage(pVideoFrame);

	cvReleaseCapture(&pCapture);	
	
	return pSnapShot;
}
	
void show(IplImage* pImage) {
	cvNamedWindow("Image Preview", CV_WINDOW_AUTOSIZE);
	cvShowImage("Image Preview", pImage);
	int key = cvWaitKey(0);
	if(key == 27) // ESC
		_break = 1;
	else if(key == 'z') // Z
		_skip = 1;	
	cvDestroyWindow("Image Preview");
}

void collect(int count, int id, int start) {
	int i = start;	
	while(i < start+count) {
		_skip = _break = 0;
		char filename[100] = { 0 };
		sprintf(filename, "./%s/%d.png", _DIR, i);
		IplImage* pImg = snapshot();	
		
		if(!pImg)
			continue;
			
		IplImage* pFace = detect(pImg);	
		if(!pFace || _skip)	// wciśnięto klawisz Z
			continue;
		
		// zapisanie zaakceptowanego zdjęcia do katalogu klasy
		if(!_break)
			cvSaveImage(filename, pFace, 0);
		
		cvReleaseImage(&pImg);
		cvReleaseImage(&pFace);
		
		if(_break) // wciśnięto klawisz ESC
			exit(0);
		++i;
	}	
}
